package demo.security;

import java.util.HashMap;
import java.util.Map;

import org.springframework.boot.autoconfigure.condition.ConditionMessage;
import org.springframework.boot.autoconfigure.condition.ConditionOutcome;
import org.springframework.boot.autoconfigure.condition.SpringBootCondition;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ConditionContext;
import org.springframework.context.annotation.Conditional;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.type.AnnotatedTypeMetadata;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;
import org.springframework.security.oauth2.provider.token.store.JwtTokenStore;
import org.springframework.security.oauth2.provider.token.store.KeyStoreKeyFactory;

import demo.security.JwtTokenStoreConfiguration.JWTTokenStoreCondition;
import lombok.extern.slf4j.Slf4j;

@Slf4j
@Configuration
@Conditional(JWTTokenStoreCondition.class)
public class JwtTokenStoreConfiguration {
	
	@Bean("JwtTokenStore")
	@Primary
	public TokenStore tokenStore() {
		return new JwtTokenStore(accessTokenConverter());
	}
	
	@Bean("JwtTokenEnhancer")
	@Primary
	public TokenEnhancer tokenEnhancer() {
	    return accessTokenConverter();
	}
	
	@Bean("JwtAccessTokenConverter")
	public JwtAccessTokenConverter accessTokenConverter() {
		
		JwtAccessTokenConverter converter = new JwtAccessTokenConverter() {
			@Override
			public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
				
				log.debug("*********JwtAccessTokenConverter:enhance accessToken:{}",accessToken);
				
				Map<String, Object> customInfo = new HashMap();

				customInfo.put("organization", "demo-jwt");
				
				//customInfo.put("principal", authentication.getPrincipal());
				
				((DefaultOAuth2AccessToken)accessToken).setAdditionalInformation(customInfo);
							
				return super.enhance(accessToken, authentication);
			}			
		};
		
		KeyStoreKeyFactory factory = new KeyStoreKeyFactory(new ClassPathResource("demo-jwt.jks"),"changeit".toCharArray());
		
		converter.setKeyPair(factory.getKeyPair("demo-jwt"));
		
		return converter;
	}
	

	static class JWTTokenStoreCondition extends SpringBootCondition{

		@Override
		public ConditionOutcome getMatchOutcome(ConditionContext context, AnnotatedTypeMetadata metadata) {
	
			ConditionMessage.Builder message = ConditionMessage.forCondition("demo.tokenStore Condition");
			
			String tokenStore = context.getEnvironment().getProperty("demo.tokenStore");
			if("JWT".equals(tokenStore)) {
				return ConditionOutcome.match(message.foundExactly("JWT tokenStore"));
			}
			return ConditionOutcome.noMatch(message.didNotFind("JWT tokenStore").atAll());
		}	
	}
}

